from langchain_community.chat_models import ChatOpenAI
from langchain.agents import initialize_agent, AgentType
from langchain_core.tools import StructuredTool  # 改用结构化工具
from pydantic import BaseModel, Field  # 用于定义参数模型

# 1. 定义工具参数模型（明确输入参数的类型和说明）
class AddInput(BaseModel):
    a: int = Field(description="第一个整数")
    b: int = Field(description="第二个整数")

class SubtractInput(BaseModel):
    a: int = Field(description="被减数（整数）")
    b: int = Field(description="减数（整数）")

# 2. 定义工具函数（保持不变）
def add(a: int, b: int) -> int:
    """返回两个数的和"""
    return a + b

def subtract(a: int, b: int) -> int:
    """返回两个数的差（a - b）"""
    return a - b

# 3. 用 StructuredTool 封装工具（指定参数模型）
tools = [
    StructuredTool(
        name="Addition",
        func=add,
        description="计算两个整数的和",
        args_schema=AddInput  # 绑定参数模型
    ),
    StructuredTool(
        name="Subtraction",
        func=subtract,
        description="计算两个整数的差（a - b）",
        args_schema=SubtractInput  # 绑定参数模型
    )
]

# 4. 配置本地LLM（不变）
llm = ChatOpenAI(
    model_name="DeepSeekR1",
    openai_api_base="http://168.17.100.17:1025/v1",
    openai_api_key="dummy"
)

# 5. 初始化Agent（不变）
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True
)

# 6. 测试
if __name__ == "__main__":
    print(agent.run("计算3加5的结果"))  # 输出：8
    print(agent.run("10减去4等于多少？"))  # 输出：6
